import mitmproxy.http
from mitmproxy import http, options
from mitmproxy.tools.dump import DumpMaster

from _0log import logger

opts = options.Options(console_log_level="warning")
m = DumpMaster(opts)
m.run()


# 需求是这样的：
#
# 因为百度搜索是不靠谱的，所有当客户端发起百度搜索时，记录下用户的搜索词，再修改请求，将搜索词改为“360 搜索”；
# 因为 360 搜索还是不靠谱的，所有当客户端访问 360 搜索时，将页面中所有“搜索”字样改为“请使用谷歌”。
# 因为谷歌是个不存在的网站，所有就不要浪费时间去尝试连接服务端了，所有当发现客户端试图访问谷歌时，直接断开连接。
# 将上述功能组装成名为 Joker 的 addon，并保留之前展示名为 Counter 的 addon，都加载进 mitmproxy。
class ListenHttp:

    def http_connect(self, flow: mitmproxy.http.HTTPFlow):
        host = flow.request.host
        logger.info("http_connect host : %s" % host)
        # 确认客户端是想访问 www.google.com
        if host == "www.google.com":
            # 返回一个非 2xx 响应断开连接
            flow.response = http.HTTPResponse.make(404)

    def request(self, flow: mitmproxy.http.HTTPFlow):
        # 忽略非百度搜索地址
        if flow.request.host != "www.baidu.com" or not flow.request.path.startswith("/s"):
            return

        # 确认请求参数中有搜索词
        if "wd" not in flow.request.query.keys():
            logger.info("can not get search word from %s" % flow.request.pretty_url)
            return

        # 输出原始的搜索词
        logger.info("catch search word: %s" % flow.request.query.get("wd"))
        # 替换搜索词为“360搜索”
        flow.request.query.set_all("wd", ["360搜索"])

    def response(self, flow: mitmproxy.http.HTTPFlow):
        # 忽略非 360 搜索地址
        if flow.request.host != "www.so.com":
            return

        # 将响应中所有“搜索”替换为“请使用谷歌”
        text = flow.response.get_text()
        text = text.replace("搜索", "请使用谷歌")
        flow.response.set_text(text)

    def error(self, flow: mitmproxy.http.HTTPFlow):
        pass


addons = [
    ListenHttp()
]
